import sys
import logging
from tqdm import tqdm
import numpy as np
from sklearn.datasets import load_digits
from sklearn.model_selection import train_test_split
from pinecone import Pinecone, ServerlessSpec

# 配置logging，包含日期
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s %(levelname)s: %(message)s",
    datefmt="%Y-%m-%d %H:%M:%S"
)

# 初始化 Pinecone
pinecone = Pinecone(api_key="pcsk_6aQeZm_5UcBB6iZSR4j5gnmoa1UZm2oJdmWuQoTUQqkDiwN7nApxibHDX2VC6Nb1yp2zZP")
index_name = "mnist-train-index"

# 删除已存在的索引
existing_indexes = pinecone.list_indexes()
if any(index['name'] == index_name for index in existing_indexes):
    pinecone.delete_index(index_name)

# 创建新索引
pinecone.create_index(
    name=index_name,
    dimension=64,
    metric="euclidean",
    spec=ServerlessSpec(cloud="aws", region="us-east-1")
)
index = pinecone.Index(index_name)

# 等待索引 ready
import time
while True:
    desc = pinecone.describe_index(index_name)
    status = getattr(desc.status, "ready", None)
    if status is True:
        break
    time.sleep(1)

# 加载数据并划分
digits = load_digits(n_class=10)
X = digits.data
y = digits.target
X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.2, random_state=42, stratify=y
)

# 上传训练数据到 Pinecone，带进度条
vectors = []
for i in range(len(X_train)):
    vectors.append({
        "id": str(i),
        "values": X_train[i].tolist(),
        "metadata": {"label": int(y_train[i])}
    })

batch_size = 1000
logging.info("开始上传训练数据到 Pinecone ...")
for i in tqdm(range(0, len(vectors), batch_size), desc="上传进度"):
    batch = vectors[i:i + batch_size]
    index.upsert(batch, namespace="")
logging.info(f"成功创建索引，并上传了{len(X_train)}条数据")

# 检查索引数据量
stats = index.describe_index_stats()
logging.info(f"索引统计信息：{stats}")

# 用测试集评估准确率，带进度条
correct = 0
k = 11
logging.info("开始用 Pinecone 检索测试集 ...")
for i in tqdm(range(len(X_test)), desc="测试进度"):
    query_vec = X_test[i].tolist()
    results = index.query(
        vector=query_vec,
        top_k=k,
        include_metadata=True,
        namespace=""
    )
    labels = [match['metadata']['label'] for match in results['matches']]
    if labels:
        pred = max(set(labels), key=labels.count)
        if pred == int(y_test[i]):
            correct += 1

accuracy = correct / len(X_test)
logging.info(f"当k=11时，使用Pinecone的准确率为：{accuracy:.4f}")

# 可选：删除索引（如不需要保留）
# pinecone.delete_index(index_name)
